package main

import (
	"fmt"
	"image"
	"image/color"
	"image/png"
	"log"
	"os"
	"time"
)

func openPNG(filename string) (image.Image, error) {
	// TODO: check filename end with .png
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	img, err := png.Decode(file)
	return img, err
}

func savePNG(img image.Image, filename string) error {
	// TODO: check filename end with .png
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	if err = png.Encode(file, img); err != nil {
		return err
	}
	return nil
}

func normalizeEnergy2Image(energy [][]float32) image.Image {
	if len(energy) == 0 {
		fmt.Fprintln(os.Stderr, "empty input")
		os.Exit(-1)
	}
	maxVal := float32(-1)
	rows, cols := len(energy), len(energy[0])
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			if maxVal < energy[i][j] {
				maxVal = energy[i][j]
			}
		}
	}

	img := image.NewGray(image.Rect(0, 0, rows, cols))
	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			var c uint8
			c = uint8(energy[i][j] / maxVal * 255)
			img.SetGray(i, j, color.Gray{c})
		}
	}
	return img
}

func computeEnergy(img image.Image) ([][]float32, error) {
	_Y, _X := img.Bounds().Max.Y, img.Bounds().Max.X
	energy := make([][]float32, _X)
	for i := 0; i < _X; i++ {
		energy[i] = make([]float32, _Y)
	}

	RGBASub := func(c1, c2 color.RGBA) color.RGBA {
		c := color.RGBA{0, 0, 0, 255}
		c.R = c1.R - c2.R
		c.G = c1.G - c2.G
		c.B = c1.B - c2.B
		return c
	}
//	tic := time.Now()
	for i := 0; i < _X; i++ {
		for j := 0; j < _Y; j++ {
			cx := color.RGBA{0, 0, 0, 255}
			cy := cx
			if i != 0 && i != _X-1 {
				cx = RGBASub(img.At(i+1, j).(color.RGBA), img.At(i-1, j).(color.RGBA))
			}
			if j != 0 && j != _Y-1 {
				cy = RGBASub(img.At(i, j+1).(color.RGBA), img.At(i, j-1).(color.RGBA))
			}
			energy[i][j] += (float32(cx.R) * float32(cx.R))
			energy[i][j] += (float32(cx.G) * float32(cx.G))
			energy[i][j] += (float32(cx.B) * float32(cx.B))

			energy[i][j] += (float32(cy.R) * float32(cy.R))
			energy[i][j] += (float32(cy.G) * float32(cy.G))
			energy[i][j] += (float32(cy.B) * float32(cy.B))
		}
	}
//	fmt.Println("energy computation time usage : ", time.Since(tic))
	return energy, nil
}

// using the input energy to select a vertical seam,
// and dump the indexs of the seam to output index bins
func selectSeamVertical(energy [][]float32) []int {
	if len(energy) == 0 {
		fmt.Fprintln(os.Stderr, "empty input")
		os.Exit(-1)
	}
	_X, _Y := len(energy), len(energy[0])
	path := make([][]int, _X)
	for i := 0; i < _X; i++ {
		path[i] = make([]int, _Y)
	}
	for row := 1; row < _Y; row++ {
		for col := 0; col < _X; col++ {
			minVal := energy[col][row-1]
			path[col][row] = col
			if col > 0 && minVal > energy[col-1][row-1] {
				minVal = energy[col-1][row-1]
				path[col][row] = col - 1
			}
			if col < _X-1 && minVal > energy[col+1][row-1] {
				minVal = energy[col+1][row-1]
				path[col][row] = col + 1
			}
			energy[col][row] += minVal
		}
	}
	index := 0
	for i := 1; i < _X; i++ {
		if energy[index][_Y-1] > energy[i][_Y-1] {
			index = i
		}
	}
	outPath := make([]int, _Y)
	outPath[_Y-1] = index
	for i := _Y - 2; i >= 0; i-- { // from bottom goes up
		outPath[i] = path[outPath[i+1]][i]

	}
	return outPath
}

func selectSeamHorizontal() {}

func main() {
	tic := time.Now()
	img, err := openPNG("in.png")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println(time.Since(tic))

	for i := 0; i < 50; i++ { // 在水平方向减少100像素宽度
		fmt.Println("i = ", i)
		energy, _ := computeEnergy(img)
		// visualizeImage(energy, "energy.png")	// debug only
		path := selectSeamVertical(energy)

		_Y, _X := img.Bounds().Max.Y, img.Bounds().Max.X
		outimg := image.NewRGBA(image.Rect(0, 0, _X-1, _Y))
		for r := 0; r < _Y; r++ {
			for c := 0; c < _X; c++ {
				if path[r] < c {
					outimg.SetRGBA(c-1, r, img.At(c, r).(color.RGBA))
				} else if path[r] > c {
					outimg.SetRGBA(c, r, img.At(c, r).(color.RGBA))
				}
			}
		}
		img = outimg
	}
	savePNG(img, "seam.png")
}
